import torch
import torchvision.transforms.functional as F
import random


class RandomImageRotator:
    def __init__(self, angles=[0, 90, 180, 270]):
        self.angles = angles

    def __call__(self, img: torch.Tensor) -> torch.Tensor:
        angle = random.choice(self.angles)
        direction = random.choice(['clockwise', 'counterclockwise'])
        if direction == 'counterclockwise':
            angle = -angle
        return F.rotate(img, angle)
